import os
import cv2
import numpy as np
import shutil
import imgaug.augmenters as iaa

"""
数据集增强
1、在阈值范围内，随机调节图片亮度
2、在阈值范围内，随机调节图片动态模糊程度，角度固定为横向模糊
"""

############需要修改###################
# 输入和输出文件夹路径
input_folder = "/home/hw/dataset/tool/img"
output_folder = "/home/hw/dataset/tool/img_enhanced"
label_folder = "label"
output_label_folder = "label_enhanced"

# 各项阈值
min_brightness = 4.3  # 最小亮度
max_brightness = 4.4  # 最大亮度
min_kernel_size = 3  # 最小模糊程度
max_kernel_size = 7  # 最大模糊程度
#####################################


def random_brightness(image):
    """随机亮度调整"""
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    brightness = np.random.uniform(min_brightness, max_brightness)
    hsv[:, :, 2] = np.clip(hsv[:, :, 2] * brightness, 0, 255).astype(np.uint8)
    return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)


def apply_motion_blur(image, max_kernel_size=7):
    """随机运动模糊"""
    kernel_size = np.random.randint(3, max_kernel_size)
    if kernel_size % 2 == 0:  # 确保内核大小为奇数
        kernel_size += 1
    seq = iaa.Sequential(
        [iaa.MotionBlur(k=kernel_size, angle=(78, 102))]  ##### 此处可以调节角度
    )
    return seq.augment_image(image)


# 创建输出文件夹
os.makedirs(output_folder, exist_ok=True)
# os.makedirs(output_label_folder, exist_ok=True)

# 处理每张图像
for filename in os.listdir(input_folder):
    if filename.endswith((".jpg", ".jpeg", ".png")):  # 仅处理图像文件
        # 读取图像
        img_path = os.path.join(input_folder, filename)
        image = cv2.imread(img_path)

        # 应用随机亮度调整
        image_enhanced = random_brightness(image)

        # 应用随机运动模糊
        # image_enhanced = apply_motion_blur(image_enhanced)

        # 保存增强后的图像
        output_path = os.path.join(output_folder, filename)
        print(img_path,"->",output_path)

        cv2.imwrite(output_path, image_enhanced)

        # 复制对应的txt文件
        # label_file_path = os.path.join(
        #     label_folder, os.path.splitext(filename)[0] + ".txt"
        # )
        # output_label_file_path = os.path.join(
        #     output_label_folder, os.path.splitext(filename)[0] + ".txt"
        # )
        # shutil.copy(label_file_path, output_label_file_path)

print("数据增强完成！")
